package com.electromagnetic.industry.software.manage.service.serviceimpl;

import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.text.StrFormatter;
import cn.hutool.core.util.StrUtil;
import cn.hutool.crypto.digest.DigestUtil;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.electromagnetic.industry.software.common.pojo.PageFile;
import com.electromagnetic.industry.software.common.pojo.UserLoginInfo;
import com.electromagnetic.industry.software.common.resp.ElectromagneticResult;
import com.electromagnetic.industry.software.common.util.ElectromagneticResultUtil;
import com.electromagnetic.industry.software.common.util.EleIdCreator;
import com.electromagnetic.industry.software.common.util.OfficeFileUtil;
import com.electromagnetic.industry.software.common.util.UserThreadLocal;
import com.electromagnetic.industry.software.manage.config.ElePropertyConfig;
import com.electromagnetic.industry.software.manage.mapper.AiFileUploadRecordMapper;
import com.electromagnetic.industry.software.manage.mapper.AiQuestionRecordMapper;
import com.electromagnetic.industry.software.manage.pojo.models.AiFileUploadRecord;
import com.electromagnetic.industry.software.manage.pojo.models.AiQuestionRecord;
import com.electromagnetic.industry.software.manage.pojo.req.QueryDTO;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.document.Document;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.multipart.MultipartFile;
import reactor.core.publisher.Flux;

import java.io.File;
import java.nio.charset.Charset;
import java.util.*;

@Slf4j
@Service
public class ChatService {

    @Resource
    private VectorStore vectorStore;
    @Resource
    private OllamaChatModel model;
    //    @Resource
//    private MessageChatMemoryAdvisor messageChatMemoryAdvisor;
    @Resource
    private AiFileUploadRecordMapper aiFileUploadRecordMapper;
    @Resource
    private AiQuestionRecordMapper aiQuestionRecordMapper;
    @Resource
    private ElePropertyConfig elePropertyConfig;

//    public ElectromagneticResult<?> deleteDocument(List<String> ids) {
//        List<String> vectorIds = aiFileUploadRecordMapper.selectList(Wrappers.lambdaQuery(AiFileUploadRecord.class)
//                        .in(AiFileUploadRecord::getId, ids))
//                .stream()
//                .map(AiFileUploadRecord::getVectorId)
//                .toList();
//        vectorStore.delete(vectorIds);
//        aiFileUploadRecordMapper.update(Wrappers.lambdaUpdate(AiFileUploadRecord.class)
//                .in(AiFileUploadRecord::getId, ids)
//                .set(AiFileUploadRecord::getEffectFlag, EffectFlagEnum.NOT_EFFECTIVE.code));
//        return ElectromagneticResultUtil.success(true);
//    }

//    @Transactional(rollbackFor = Exception.class)
//    public ElectromagneticResult<?> addFromUpload(MultipartFile file) throws Exception {
//        String s = FileUtil.readString("D:\\datas\\aidocs\\HB8451民用飞机燃油系统电搭接要求.txt", Charset.defaultCharset());
//        Document document = new Document(s);
//        List<Document> newDocuments = new TokenTextSplitter().apply(List.of(document));
//        vectorStore.write(newDocuments);
//        return ElectromagneticResultUtil.success(true);
//    }


    @Transactional(rollbackFor = Exception.class)
    public ElectromagneticResult<?> addFromUpload(MultipartFile file) throws Exception {

        // 文件是否为空
        if (file.isEmpty()) {
            return ElectromagneticResultUtil.fail("-1", "文件为空");
        }

        String fileType = FileUtil.extName(file.getOriginalFilename());
        if (!elePropertyConfig.getAiAllowUploadFileTypes().contains(fileType)) {
            return ElectromagneticResultUtil.fail("-1", StrFormatter.format("当前格式 {} 不支持", fileType));
        }

        // 通过md5值判断文件是否被上传过
        String fileMd5 = DigestUtil.md5Hex(file.getInputStream());
        Long count = aiFileUploadRecordMapper.selectCount(Wrappers.lambdaQuery(AiFileUploadRecord.class)
                .eq(AiFileUploadRecord::getFileMd5, fileMd5));
        if (count > 0) {
            return ElectromagneticResultUtil.success(fileMd5);
        }
        List<Document> documents = new ArrayList<>();
        switch (Objects.requireNonNull(fileType)) {
            case "txt", "csv", "text" -> {
                String content = IoUtil.read(file.getInputStream(), Charset.defaultCharset());
                documents.add(new Document(content));
            }
            case "xls", "xlsx" -> {
                String filePath = elePropertyConfig.getEleTmpPath() + File.separator + fileMd5 + "." + fileType;
                FileUtil.writeFromStream(file.getInputStream(), filePath);
                List<PageFile> pageInfo = StrUtil.equals("xlsx", filePath) ?
                        OfficeFileUtil.parseXlsxByPage(filePath) :
                        OfficeFileUtil.parseXlsByPage(filePath);
                for (PageFile pageFile : pageInfo) {
                    Map<String, Object> metaData = Map.of("page_number", pageFile.getPageNumber(), "file_name", file.getOriginalFilename());
                    Document document = new Document(pageFile.getContent(), metaData);
                    documents.add(document);
                }
                FileUtil.del(filePath);
            }
            case "doc", "docx", "ppt", "pptx", "pdf" -> {
                List<PageFile> pageFiles;
                String pdfPath = elePropertyConfig.getEleTmpPath() + File.separator + fileMd5 + ".pdf";
                String srcPath = elePropertyConfig.getEleTmpPath() + File.separator + fileMd5 + "." + fileType;
                if (StrUtil.equals(fileType, "pdf")) {
                    FileUtil.writeFromStream(file.getInputStream(), pdfPath);
                } else {
                    FileUtil.writeFromStream(file.getInputStream(), srcPath);
                    if (Arrays.asList("doc", "docx").contains(fileType)) {
                        OfficeFileUtil.doc2pdf(srcPath, pdfPath);
                    } else {
                        OfficeFileUtil.ppt2pdf(srcPath, pdfPath);
                    }
                }
                pageFiles = OfficeFileUtil.parsePdfByPage(pdfPath);
                StringBuilder stringBuilder = new StringBuilder();
                pageFiles.forEach(page -> stringBuilder.append(page.getContent()));
                Document document = new Document(stringBuilder.toString(), Map.of("file_name", file.getOriginalFilename()));
                documents.add(document);
                FileUtil.del(pdfPath);
                FileUtil.del(srcPath);
            }
            default -> {
                return ElectromagneticResultUtil.fail("-1", StrFormatter.format("当前格式 {} 不支持", fileType));
            }
        }
        List<Document> newDocuments = new ArrayList<>();
        StringBuilder stringBuilder = new StringBuilder();
        for (Document document : documents) {
            Map<String, Object> metadata = Map.of("fileName", file.getOriginalFilename(), "fileMd5", fileMd5);
            newDocuments.add(new Document(EleIdCreator.createSnowflakeId(), Objects.requireNonNull(document.getText()), metadata));
            stringBuilder.append(document.getText());
        }
        FileUtil.writeString(stringBuilder.toString(), elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
        newDocuments = new TokenTextSplitter().apply(newDocuments);
        newDocuments = newDocuments.subList(0, Math.min(elePropertyConfig.getAiMaxDocs(), newDocuments.size()));
        vectorStore.write(newDocuments);
        for (Document document : newDocuments) {
            aiFileUploadRecordMapper.insert(new AiFileUploadRecord()
                    .setId(EleIdCreator.createSnowflakeId())
                    .setVectorId(document.getId())
                    .setFileSize(file.getSize())
                    .setFileMd5(fileMd5)
                    .setFileName(file.getOriginalFilename()));
        }
        return ElectromagneticResultUtil.success(fileMd5);
    }

    public Flux<String> chatStreamStr(QueryDTO queryDTO) {
        AiQuestionRecord record = new AiQuestionRecord()
                .setQuestion(queryDTO.getMsg())
                .setId(EleIdCreator.createSnowflakeId());
        UserLoginInfo userLoginInfo = new UserLoginInfo();
        userLoginInfo.setUserId(queryDTO.getUserId());
        UserThreadLocal.set(userLoginInfo);
        aiQuestionRecordMapper.insert(record);
        UserThreadLocal.remove();
        List<Document> documents = vectorStore.similaritySearch(queryDTO.getMsg());
        Set<String> files = new LinkedHashSet<>();
        assert documents != null;
        for (Document document : documents) {
            files.add(document.getMetadata().get("fileMd5").toString());
        }
        StringBuilder info = new StringBuilder();
        for (String fileMd5 : files) {
            String s = FileUtil.readString(elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
            info.append(s);
        }
        String augmentedPrompt = "基于以下信息回答：" + info + "\n\n" + queryDTO.getMsg();
        return ChatClient.builder(model)
                .defaultSystem("必须用中文回答")
//                .defaultAdvisors(messageChatMemoryAdvisor)
                .defaultOptions(OllamaOptions
                        .builder()
                        .temperature(0.3)
                        .build())
                .build()
                .prompt(augmentedPrompt)
                .stream()
                .content();
    }

    public Flux<ChatResponse> chatStreamResponse(QueryDTO queryDTO) {
        AiQuestionRecord record = new AiQuestionRecord().setQuestion(queryDTO.getMsg()).setId(EleIdCreator.createSnowflakeId());
        UserLoginInfo userLoginInfo = new UserLoginInfo();
        userLoginInfo.setUserId(queryDTO.getUserId());
        UserThreadLocal.set(userLoginInfo);
        aiQuestionRecordMapper.insert(record);
        UserThreadLocal.remove();
        List<Document> documents = vectorStore.similaritySearch(SearchRequest.builder()
//                .topK(20)
//                .similarityThreshold(0.7)
                .query(queryDTO.getMsg())
                .build());

        Set<String> files = new LinkedHashSet<>();
        assert documents != null;
        for (Document document : documents) {
            files.add(document.getMetadata().get("fileMd5").toString());
        }
        StringBuilder info = new StringBuilder();
        for (String fileMd5 : files) {
            String s = FileUtil.readString(elePropertyConfig.getAiFileUploadStoreDir() + File.separator + fileMd5, Charset.defaultCharset());
            info.append(s);
        }

        String augmentedPrompt = "基于以下信息回答：" + info + "\n\n" + queryDTO.getMsg();
        return ChatClient.builder(model)
                .defaultSystem("必须用中文回答")
//                .defaultAdvisors(messageChatMemoryAdvisor)
                .defaultOptions(OllamaOptions
                        .builder()
                        .temperature(0.3)
                        .build())
                .build()
                .prompt(new Prompt(new UserMessage(augmentedPrompt)))
                .stream()
                .chatResponse();
    }
}
